import json
import os
import re
import subprocess
from datetime import datetime, timedelta

import psutil

from libs import config
from libs.db import DB
from libs.log import Log
from libs.window import AdsWindow


class TaskStatus:
    PENDING = 'pending'
    RUNNING = 'running'
    DONE = 'done'
    TERMINATED = 'terminated'
    FAIL = 'fail'


class SchedulerV2:
    def __init__(self):
        self.datetime_format = '%Y-%m-%d %H:%M:%S'
        self.logger = Log('scheduler')
        self.db = DB()
        self.db.run_sql("""
            CREATE TABLE IF NOT EXISTS task (
                id INTEGER PRIMARY KEY,
                user_id INTEGER NOT NULL,
                task_id TEXT NOT NULL,
                create_time DATETIME NOT NULL,
                execute_time DATETIME NOT NULL,
                command TEXT NOT NULL,
                priority INTEGER DEFAULT 0,
                pid TEXT NULL,
                start_time DATETIME NULL,
                end_time DATETIME NULL,
                status TEXT NOT NULL
            );
        """)

        self.db.run_sql("""
            CREATE TABLE IF NOT EXISTS running_window (
                id INTEGER PRIMARY KEY,
                user_id INTEGER NOT NULL,
                start_time DATETIME NULL
            );
        """)

    def reset(self):
        self.db.run_sql("DROP TABLE task")
        print('reset all tasks')
        self.__init__()

    @staticmethod
    def clear():
        db = DB()
        row = db.fetch_one("SELECT count(*) as row_count from task where  create_time <= DATETIME('now', '-24 hours')")
        print(row)
        db.run_sql("DELETE FROM task where create_time <= DATETIME('now', '-24 hours')")
        print('task data cleared')

    def list_running_windows(self):
        sql = """
                SELECT user_id 
                FROM running_window
                ORDER BY start_time ASC
                """
        rows = self.db.fetch_all(sql)
        result = []
        for row in rows:
            result.append(row[0])
        return result

    def save_running_window(self, user_id):
        self.db.run_sql("DELETE FROM running_window WHERE user_id = ?", (user_id,))
        sql = """
                INSERT INTO running_window(user_id, start_time)
                VALUES (?, ?)
                """
        self.db.run_sql(sql, (
            user_id,
            datetime.now().strftime(self.datetime_format)
        ))

    def close_window(self, user_id):
        try:
            AdsWindow(user_id).close()
            self.db.run_sql("DELETE FROM running_window WHERE user_id = ?", (user_id,))
        except:
            pass

    @staticmethod
    def is_process_running(pid):
        return psutil.pid_exists(int(pid))

    def is_task_pending(self, task_id):
        sql = """
        SELECT count(*) 
        FROM task
        WHERE task_id = ? and  status = ?
        """
        result = self.db.fetch_one(sql, (task_id, TaskStatus.PENDING))
        return result[0] > 0

    def is_task_running(self, task_id):
        sql = """
        SELECT count(*) 
        FROM task
        WHERE task_id = ? and  status = ?
        """
        result = self.db.fetch_one(sql, (task_id, TaskStatus.RUNNING))
        return result[0] > 0

    def save_task(self, user_id, new_id, command, priority, execute_time: datetime, status):
        sql = """
        INSERT INTO task(user_id, task_id, command, priority, execute_time, create_time, status)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """
        self.db.run_sql(sql, (
            user_id,
            new_id,
            command,
            priority,
            execute_time.strftime(self.datetime_format),
            datetime.now().strftime(self.datetime_format),
            status
        ))

    def delete_pending_task(self, script, user_id):
        sql = """
        DELETE FROM task where status = 'pending' and task_id = ?
        """
        task_id = self.get_task_id(user_id, script)
        self.db.run_sql(sql, (task_id, ))

    def add_task(self, script, user_id, args=None,
                 execute_time: datetime = None,
                 delay_seconds=0,
                 ignore_running=True,
                 priority=0
                 ):
        if args is None:
            args = []
        new_id = self.get_task_id(user_id, script)
        if execute_time is None:
            execute_time = datetime.now()
        if delay_seconds is not None:
            execute_time = execute_time + timedelta(seconds=delay_seconds)

        if self.is_task_pending(new_id):
            self.logger.debug(f"task[{new_id}] is pending...")
            return 0

        if not ignore_running and self.is_task_running(new_id):
            self.logger.debug(f"task[{new_id}] is running...")
            return 0

        command = f"python {script} --user_id={user_id}"
        if len(args) > 0:
            for key in args:
                command += f" --{key}={args.get(key)}"
        self.save_task(user_id, new_id, command, priority, execute_time, TaskStatus.PENDING)
        self.logger.debug('add_task ' + json.dumps({'command': command}))
        return 1

    def next_task(self):
        sql = """
        SELECT id, user_id, task_id, command, execute_time, status FROM task WHERE status = ? or status = ?
        ORDER BY priority desc, execute_time asc
        """
        rows = self.db.fetch_all(sql, (TaskStatus.PENDING, TaskStatus.RUNNING))
        pending_tasks = []
        running_tasks = []
        for row in rows:
            status = row[5]
            task = {
                "id": row[0],
                "user_id": row[1],
                "task_id": row[2],
                "command": row[3],
                "execute_time": row[4]
            }
            if status == TaskStatus.PENDING:
                pending_tasks.append(task)
            if status == TaskStatus.RUNNING:
                running_tasks.append(task)
        process_limit = config.get_process_limit()
        exclude_accounts = config.get_exclude_accounts()
        channel_process_counts = {}
        running_users = []
        process_limits = config.get_process_limit()
        if 'total' in process_limits:
            total_process_limit = process_limits.get('total')
        else:
            total_process_limit = sum(process_limits.values())
        if len(running_tasks) >= total_process_limit:
            self.logger.debug(f'running_tasks[{len(running_tasks)}] exceeds total_process_limit[{total_process_limit}]')
            return None
        for running_task in running_tasks:
            running_users.append(str(running_task.get('user_id')))
            channel = self.get_channel_from_command(running_task.get('command'))
            channel_process_counts[channel] = channel_process_counts.get(channel, 0) + 1
            if channel.startswith('farm_'):
                channel_process_counts['farm'] = channel_process_counts.get('farm', 0) + 1

        sleep_configs = config.get_sleep_configs()

        running_only = True  # 优先检查开启的窗口是否有任务
        running_user_ids = self.list_running_windows()
        for _ in range(2):
            for pending_task in pending_tasks:
                user_id = pending_task.get('user_id')
                if running_only and user_id not in running_user_ids:
                    continue
                # 检查账号是否在休息
                if user_id in sleep_configs:
                    sleep_config = sleep_configs.get(user_id)
                    start = sleep_config[0]
                    end = sleep_config[1]
                    hour = datetime.now().hour
                    if end > start and start <= hour <= end:
                        continue
                    if end < start and (hour >= start or hour <= end):
                        continue
                channel = self.get_channel_from_command(pending_task.get('command'))
                if int(user_id) in exclude_accounts:
                    continue
                if str(user_id) in running_users:
                    continue
                if pending_task['execute_time'] > datetime.now().strftime(self.datetime_format):
                    continue

                channel_limit = process_limit.get(channel, 1)
                channel_running = channel_process_counts.get(channel, 0)
                if channel_running >= channel_limit:
                    continue
                if channel.startswith('farm_'):
                    channel = 'farm'
                    channel_limit = process_limit.get(channel, 1)
                    channel_running = channel_process_counts.get(channel, 0)

                    if channel_running >= channel_limit:
                        continue
                return pending_task
            running_only = False

        self.logger.debug(f"no_coming_tasks")
        return None

    def run(self):
        # 读取待执行的任务列表
        task = self.next_task()
        if task is None:
            return
        task_id = task.get('id')
        command = task.get('command')
        command += f" --task_id={task_id}"
        working_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../')
        env = os.environ.copy()
        log_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                     f"../runtime/debug/debug_{task.get('user_id')}.log")
        self.save_running_window(task.get('user_id'))
        with open(log_file_path, 'a') as f:
            process = subprocess.Popen(
                command,
                shell=True,
                start_new_session=True,
                cwd=working_dir,
                env=env,
                stdout=f,
                stderr=f
            )
        self.logger.debug('run_task ' + json.dumps({'command': command}))
        sql = """
        UPDATE task SET pid = ? , start_time = ?, status = ?
        WHERE id = ?
        """
        self.db.run_sql(sql, (
            process.pid,
            datetime.now().strftime(self.datetime_format),
            TaskStatus.RUNNING,
            task_id
        ))

    def finish(self, task_id, status):
        sql = """
                UPDATE task SET end_time = ?, status = ?
                WHERE id = ?
                """
        self.db.run_sql(sql, (
            datetime.now().strftime(self.datetime_format),
            status,
            task_id
        ))

    def giveup(self, task_id):
        sql = """
        DELETE FROM task where id = ?
        """
        self.db.run_sql(sql, (task_id,))

    @staticmethod
    def get_task_id(user_id, script):
        task_name = str(script).replace('.py', '')
        return f"{user_id}_{task_name}"

    @staticmethod
    def get_channel_from_command(cmd):
        script = cmd.split(' ')[1].replace('.py', '')
        if script == 'farm':
            # 使用正则表达式匹配 soil_id
            match = re.search(r'--soil_id=(\d+)', cmd)
            soil_id = match.group(1)
            channel = f"farm_{soil_id}"
        else:
            channel = script
        return channel

    def get_user_latest_tasks(self):
        sql = """
        SELECT t.id, t.user_id, t.status, t.pid, t.start_time, t.end_time
        FROM task t
        JOIN (
            SELECT user_id, MAX(start_time) AS max_start_time
            FROM task
            WHERE start_time IS NOT NULL
            GROUP BY user_id
        ) latest_task
        ON t.user_id = latest_task.user_id AND t.start_time = latest_task.max_start_time;
        """
        rows = self.db.fetch_all(sql)
        result = {}
        for row in rows:
            task_id = int(row[0])
            user_id = int(row[1])
            status = str(row[2])
            pid = row[3]
            start_time = str(row[4])
            end_time = str(row[5])
            if status == TaskStatus.RUNNING:
                if not self.is_process_running(pid):
                    self.finish(task_id, TaskStatus.TERMINATED)
                    continue
            result[user_id] = {
                "task_id": task_id,
                "status": status,
                "start_time": start_time,
                "end_time": end_time
            }
        return result

    def get_pending_and_running_task_ids(self):
        sql = """
        SELECT task_id, status 
        FROM task
        WHERE status in ('pending', 'running')
        """
        rows = self.db.fetch_all(sql)
        result = []
        for row in rows:
            result.append(row[0])
        return result

    def next_schedule_time(self, user_id, task_id, duration=0):
        running_sql = """
        SELECT task_id
        FROM task
        WHERE status = 'running' and user_id = ?
        """
        running_rows = self.db.fetch_one(running_sql, (user_id,))
        if running_rows is not None:
            # 有运行中的任务
            return None

        done_sql = """
        SELECT end_time
        FROM task
        WHERE status = 'done' and task_id = ? and user_id =?
        ORDER BY id DESC
        """
        row = self.db.fetch_one(done_sql, (task_id, user_id))
        if row is not None:
            return datetime.strptime(row[0], self.datetime_format) + timedelta(seconds=duration)
        else:
            return datetime.now()
